import torch
from torch import nn
from config import HP
from dataset_face import invTrans
import matplotlib.pyplot as plt
import torchvision.utils as vutil


class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()

        self.projection_layer = nn.Linear(
            in_features=HP.z_dim,
            out_features=4 * 4 * 1024
        )  # 1.feature transform  2.reshape

        self.generator = nn.Sequential(
            # TransposeConv:1
            nn.ConvTranspose2d(
                in_channels=1024,
                out_channels=512,
                kernel_size=(4, 4),
                stride=(2, 2),
                padding=(1, 1),
                bias=False
            ),
            nn.BatchNorm2d(512),
            nn.ReLU(),

            # TransposeConv:2
            nn.ConvTranspose2d(
                in_channels=512,
                out_channels=256,
                kernel_size=(4, 4),
                stride=(2, 2),
                padding=(1, 1),
                bias=False
            ),
            nn.BatchNorm2d(256),
            nn.ReLU(),

            # TransposeConv:3
            nn.ConvTranspose2d(
                in_channels=256,
                out_channels=128,
                kernel_size=(4, 4),
                stride=(2, 2),
                padding=(1, 1),
                bias=False
            ),
            nn.BatchNorm2d(128),
            nn.ReLU(),

            # TransposeConv:4
            nn.ConvTranspose2d(
                in_channels=128,
                out_channels=HP.data_channels,  # output channel = 3
                kernel_size=(4, 4),
                stride=(2, 2),
                padding=(1, 1),
                bias=False
            ),
            nn.Tanh()  # [0, 1]
        )

    def forward(self, latent_z):
        z1 = self.projection_layer(latent_z)
        z_projected = z1.view(-1, 1024, 4, 4)  # [N, C, H, W]
        return self.generator(z_projected)

    @staticmethod
    def weight_init(layer):
        layer_class_name = layer.__class__.__name__
        if 'Conv' in layer_class_name:
            nn.init.normal_(layer.weight.data, 0.0, 0.02)
        elif 'BatchNorm' in layer_class_name:
            nn.init.normal_(layer.weight.data, 1.0, 0.02)  # in_place
            nn.init.normal_(layer.bias.data, 0.)


if __name__ == '__main__':
    z = torch.randn(size=(64, 100))
    G = Generator()
    g_out = G(z)  # Generator的output
    print(g_out.size())

    grid = vutil.make_grid(g_out, nrow=8)  # 将小图片设计成大图片
    plt.imshow(invTrans(grid).permute(1, 2, 0))
    plt.show()
